Skip to main content

🚀 AI Infra 大话西游实战之Transformer小白学习笔记(一)

Google 大神 × 👓 Ringi 导学版(公式全兼容 / 小白友好 / AI Infra 实战贯通) 📌 导读寄语: 学习 Transformer 不是为了死记硬背公式,而是要建立清晰的数据流动直觉与硬件执行心智模型。 本笔记以 Google 大神的第一性原理视角,结合极客导师 Ringi 的手办级 3D 架构工坊插图,把每个公式、每个张量形状(Shape)以及后续在 CUDA / 分布式训练 / vLLM 推理中的对应关系彻底讲透!

📑 目录导航


1. 为什么 AI Infra 工程师必须懂 Transformer

1.1 汽车发动机比喻

如果你准备深入 AI Infra(AI 基础设施)——无论是在底层写 CUDA C++ 算子、在集群搞 Megatron-LM / DeepSpeed 分布式训练,还是在 vLLM / TensorRT-LLM 搞高性能推理服务——你天天面对的核心工作对象都是同一个东西:Transformer 模型。 这就像汽车赛车工程师必须懂发动机结构一样:
  • 你不需要自己从头画机械草图(那是算法研究员 / 架构师的活);
  • 但你必须清楚知道曲轴在哪里、活塞怎么往复运动、喷油嘴怎么供油——否则你根本不知道哪个环节卡住了吞吐,该去优化哪个部件!

1.2 AI Infra 各层级与 Transformer 模块映射全景

Ringi 导师解构:AI Infra 发动机与 Transformer 模块映射全景工坊 AI Infra 的每一项核心优化技术,都直接精确对应 Transformer 的某一个数学算子或物理结构:

映射关系速查表

💡 Ringi 划重点: “不懂 Transformer,就不知道自己在优化什么。” 后续在工程中遇到的每一个优化算法,都能在本文中找到它作用的具象代码行与数学公式。

2. Transformer 网络结构全貌

在深入每个矩阵乘法之前,我们先站在 Google 架构师的高度俯瞰整体蓝图。

2.1 原始 Transformer:Encoder-Decoder 架构

2017 年 Google 经典论文 《Attention Is All You Need》 提出的原始结构由两部分组成:
  • Encoder:全向双向注意力(每个 token 都可以看到整句话的所有 token),主要负责阅读理解。
  • Decoder:带因果掩码的注意力(不能看未来的词)+ 交叉注意力(Cross-Attention,向 Encoder 索要信息),负责自回归生成。

2.2 当前大模型王者:Decoder-only 架构

GPT 系列(以及后来的 LLaMA、Mistral、Qwen、DeepSeek 等)证明了一个震撼业界的结论: 根本不需要 Encoder 和 Cross-Attention,只用 Decoder 就足够解决人类所有语言任务! 当前主流 LLM 结构采用极度优雅的三段式: Ringi 导师解构:Transformer 整体三段式架构全景

经典三段式数据流动流程:

  1. 输入阶段(Input Stage):
Prompt Tokens⟶Token Embedding (词表查找)⟶注入位置编码 (如 RoPE)\text{Prompt Tokens} \longrightarrow \text{Token Embedding (词表查找)} \longrightarrow \text{注入位置编码 (如 RoPE)}
  1. 核心堆叠阶段(Repeated N Blocks):
Xl+1=Xl+Self-Attention(Norm(Xl))+FFN(Norm(… ))X_{l+1} = X_l + \text{Self-Attention}(\text{Norm}(X_l)) + \text{FFN}(\text{Norm}(\dots))
  1. 输出阶段(Output Stage):
Final Norm⟶LM Head 线性映射⟶Softmax 采样⟶预测下一个 Token\text{Final Norm} \longrightarrow \text{LM Head 线性映射} \longrightarrow \text{Softmax 采样} \longrightarrow \text{预测下一个 Token}

2.3 三大架构变体对比与演进逻辑

💡 Ringi 导师解惑: 为什么 Decoder-only 能统治整个大模型时代?
  1. 万物皆可 Next-Token Prediction:翻译、推理、写代码、甚至图像多模态,全部统一为自回归预测,不需要为不同任务换模型架构。
  2. 工程部署极致高效:全模型只有一种标准的 Decoder Block,显存分配、KV Cache 调度、张量并行通信模式整齐划一。

3. Self-Attention 机制(核心引擎)

Self-Attention 是 Transformer 的心脏,也是算力与显存消耗最密集的重灾区。

3.1 直觉理解:图书馆查资料与信息加权聚合

算谁重要(注意力权重)⟶按重要程度加权拿取信息\boxed{\text{算谁重要(注意力权重)} \longrightarrow \text{按重要程度加权拿取信息}} 先建立一个小白直觉:Attention 就是让句子里的每个词去“关注”其他所有词,然后按关注度加权打包信息。 举个例子:句子为 小明 喜欢 吃 苹果。 当模型处理 吃 这个动作时:
  • 它需要知道“谁在吃”(小明,占 10% 注意力);
  • “吃什么”(苹果,占 70% 注意力);
  • 最终 吃 这个词更新后的特征向量为:
New-Feature吃=0.1×小明+0.1×喜欢+0.1×吃+0.7×苹果\text{New-Feature}_{\text{吃}} = 0.1 \times \text{小明} + 0.1 \times \text{喜欢} + 0.1 \times \text{吃} + 0.7 \times \text{苹果}

3.2 Q、K、V 的物理意义:同一个 Token 的三副眼镜

在 Self-Attention 中,每个 token 都会通过三组不同的投影矩阵,同时拥有三个身份: Ringi 导师解构:Self-Attention 机制三工位流水线
  • Query (QQ):提问者,发出检索需求(“我是‘吃’,我在找什么宾语?”);
  • Key (KK):索引标牌,用于和别人的 Query 匹配(“我是‘苹果’,我是可食用水果”);
  • Value (VV):真正的内容载荷,匹配成功后贡献给别人(“提供苹果的完整语义特征”)。
📌 核心口诀: Q 与 K 点积决定【关注比例】,V 决定【真正带走什么信息】\boxed{Q \text{ 与 } K \text{ 点积决定【关注比例】,} V \text{ 决定【真正带走什么信息】}}

3.3 逐步拆解完整计算流程(6 步推导 + 数值小算盘)

假设输入序列有 NN 个 token,隐藏层维度为 dd(即输入矩阵 X∈RN×dX \in \mathbb{R}^{N \times d} )。

步骤 1:线性投影生成 Q,K,VQ, K, V

Q=XWQ,K=XWK,V=XWVQ = X W_Q, \quad K = X W_K, \quad V = X W_V 其中 X∈RN×dX \in \mathbb{R}^{N \times d}, WQ,WK,WV∈Rd×dW_Q, W_K, W_V \in \mathbb{R}^{d \times d},输出的 Q,K,VQ, K, V 形状均为 (N,d)(N, d)。
🛠️ AI Infra 视点:这是 3 次标准的 GEMM(通用矩阵乘法)操作,Tensor Core 的绝对主场。

步骤 2:计算注意力原始分数( QKTQ K^T )

衡量每对 token 之间的相关程度: S=QKT∈RN×NS = Q K^T \in \mathbb{R}^{N \times N} S[i][j]S[i][j] 表示第 ii 个 token 对第 jj 个 token 的原始打分。 小白极简数值推导: 若 q=[1,2]q = [1, 2], k=[3,4]k = [3, 4],则: q⋅kT=1×3+2×4=11q \cdot k^T = 1 \times 3 + 2 \times 4 = 11 点积数值越大,说明两个向量的方向越接近、语义越相关。

步骤 3:缩放(Scale)——为什么要除以 dk\sqrt{d_k}?

Sscaled=QKTdkS_{\text{scaled}} = \frac{Q K^T}{\sqrt{d_k}}
👓 Ringi 划重点(公式兼容与原理解析): 当维度 dkd_k 很大时(例如 dk=128d_k = 128 ), QQ 和 KK 的点积相当于 128 个独立分量相乘求和,方差会放大到 dkd_k。 这会导致点积数值极其巨大(比如上百),送入 Softmax 后输出会被“推向饱和区”——最大值变成 1,其他全变成 0,梯度几乎完全消失(Gradient Vanishing)! 除以 dk\sqrt{d_k} 就像给分数**“装上降温空调”**,把方差拉回到 1,保证反向传播梯度通畅。

步骤 4:Softmax 归一化(原始分数变注意力比例)

A=softmax(Sscaled)∈RN×NA = \text{softmax}(S_{\text{scaled}}) \in \mathbb{R}^{N \times N} 对 SscaledS_{\text{scaled}} 矩阵的每一行进行 Softmax,使每行所有元素变成正数且和为 1: A[i][j]=eSscaled[i][j]∑k=1NeSscaled[i][k]A[i][j] = \frac{e^{S_{\text{scaled}}[i][j]}}{\sum_{k=1}^{N} e^{S_{\text{scaled}}[i][k]}}

步骤 5:加权求和(用注意力比例捞取 Value)

Output=A⋅V∈RN×d\text{Output} = A \cdot V \in \mathbb{R}^{N \times d} 第 ii 个 token 的新向量就是序列中所有 token 的 VV 按照第 ii 行权重 A[i]A[i] 进行加权融合。

步骤 6:输出线性投影

Final=Output⋅WO∈RN×d\text{Final} = \text{Output} \cdot W_O \in \mathbb{R}^{N \times d} 经过权重矩阵 WO∈Rd×dW_O \in \mathbb{R}^{d \times d} 映射,完成当前多头特征的综合整理。

🏆 完整标准公式总结

Attention(Q,K,V)=softmax(QKTdk+Mask)V\boxed{\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} + \text{Mask}\right) V}

3.4 为什么复杂度是 O(N2)O(N^2)?长上下文之痛

从数学公式即可清晰看出显存与计算瓶颈:
  1. 计算量: O(N2⋅d)O(N^2 \cdot d),因为生成 N×NN \times N 个元素,每个元素需要 dd 次乘加。
  2. 显存占用: O(N2)O(N^2),必须在显存中开辟 N×NN \times N 大小的临时矩阵存储 Softmax 前后的权重。
💥 算力与显存暴击: 当上下文长度 NN 从 2K 扩展到 128K 时: (128K2K)2=642=4096 倍\left(\frac{128\text{K}}{2\text{K}}\right)^2 = 64^2 = 4096 \text{ 倍} 计算量与注意力矩阵显存暴涨 4096 倍!

3.5 Multi-Head Attention:为什么要多头?

单头 Attention 就像一个人只有单重视角(可能只关注了语法搭配)。 Multi-Head Attention(MHA) 将隐藏维度 dmodeld_{\text{model}} 切分成 hh 个独立的“头”(Head),每个头维度 dk=dmodel/hd_k = d_{\text{model}} / h。 Ringi 导师解构:Multi-Head Attention 与 GQA 分组共享机制

多头机制与张量并行(Tensor Parallelism)天然契合

在 Megatron-LM 分布式训练中:
  • 32 个头可以均分给 4 张 GPU(每张卡算 8 个头);
  • 各卡计算自身对应的 WQ,WK,WVW_Q, W_K, W_V 局部矩阵;
  • 最后各卡做一次 AllReduce 汇总,通信成本极低!

3.6 MHA →\to MQA →\to GQA 的演进直觉

为了在推理时大幅削减 KV Cache 显存占用,行业经历了三代演进:
  • MHA (Multi-Head Attention): Q,K,VQ, K, V 头部数量完全相等(1:1:1),效果好但推理显存巨大。
  • MQA (Multi-Query Attention):所有 QQ 头共享同一对 K,VK, V,KV Cache 暴降为 1/h1/h,但模型表达力有所损耗。
  • GQA (Grouped-Query Attention):折中方案!例如 8 个 QQ 头分为 2 组,每组共享 1 对 K,VK, V(LLaMA-2-70B、LLaMA-3、Mistral 标配)。

4. 前馈网络(FFN)——深度加工厂

如果说 Attention 是全班同学坐在一起开圆桌会议交换观点,那么 FFN 就是会后每个人回到自己工位上独立思考、消化记忆并形成最终结论。

4.1 角色分工:Attention 负责开会,FFN 负责独立吸收

  • Attention:发生 Token 与 Token 之间的横向信息交互(混合不同位置的信息)。
  • FFN:对 每个 Token 的特征向量做独立的纵向非线性深度加工,Token 之间互不干扰(Point-wise)。

4.2 结构解析:为什么先升维再降维?

标准 FFN 结构为两层全连接: FFN(x)=W2⋅activation(W1x+b1)+b2\text{FFN}(x) = W_2 \cdot \text{activation}(W_1 x + b_1) + b_2 Ringi 导师解构:FFN 前馈网络与 SwiGLU 升降维与门控加工厂
💡 Google 大神思考: 为什么要先升维(扩大 4 倍或 8/3 倍)再降维? 就像解压缩文件一样:高维空间有更大的特征容量,可以让模型在一个更广阔的几何空间里把交织缠绕的语义概念非线性“展开”,加工提炼后再压缩回原始维度。

4.3 参数量分析:为什么 FFN 占了模型 2/3 的参数?

我们手算一个标准的 Transformer Block:
  • 隐藏维度为 dmodeld_{\text{model}},FFN 中间维度 dff≈4dmodeld_{\text{ff}} \approx 4 d_{\text{model}}。
📌 关键结论: 在整个 Transformer 中,FFN 占据了整整约 67%(2/3)的参数量! 这也是为什么在大模型分布式并行中,FFN 的切分效率直接决定了集群吞吐;在 MoE(混合专家模型)中,正是把 FFN 复制多份切分成各个专家网络。

4.4 激活函数演进:ReLU →\to GELU →\to SwiGLU

1. ReLU

ReLU(x)=max⁡(0,x)\text{ReLU}(x) = \max(0, x) 缺点:当输入为负时导数为 0,神经元容易“永久死亡”。

2. GELU(高斯误差线性单元)

GELU(x)=x⋅Φ(x)=x⋅P(X≤x),X∼N(0,1)\text{GELU}(x) = x \cdot \Phi(x) = x \cdot P(X \le x), \quad X \sim \mathcal{N}(0, 1) 不是生硬地开/关,而是根据输入大小赋予平滑的通过概率,GPT-2/3、BERT 广泛采用。

3. SwiGLU(现代大模型标配)

SwiGLU(x)=(Swish(xWgate)⊙(xWup))Wdown\text{SwiGLU}(x) = \left( \text{Swish}(x W_{\text{gate}}) \odot (x W_{\text{up}}) \right) W_{\text{down}} 其中 Swish(z)=z⋅σ(z)\text{Swish}(z) = z \cdot \sigma(z), ⊙\odot 为逐元素相乘。
🌟 优势:引入了一个专门负责**“把关”**的门控分支 WgateW_{\text{gate}},动态控制每个特征通道的放行比例,表达能力远超单个激活函数。

5. 位置编码(Position Encoding)——赋予序列时空感

5.1 为什么 Transformer 天生是个“词序脸盲”?

回顾 Attention 计算公式: S=QKTS = Q K^T。 如果把句子 猫 抓 鼠 打乱成 鼠 抓 猫,只要输入向量跟着换行,计算出来的每对词之间的注意力值完全一样! 数学上称为 排列等变性(Permutation Equivariance)——如果不显式注入位置信息,Transformer 根本分不清主语和宾语!

5.2 经典方案:Sinusoidal 正余弦位置编码

2017 年原始 Transformer 采用基于不同频率正弦/余弦的固定绝对位置编码: PE(pos,2i)=sin⁡(pos100002i/dmodel)PE(pos,2i+1)=cos⁡(pos100002i/dmodel)\begin{aligned} \text{PE}_{(\text{pos}, 2i)} &= \sin\left(\frac{\text{pos}}{10000^{2i / d_{\text{model}}}}\right) \\ \text{PE}_{(\text{pos}, 2i+1)} &= \cos\left(\frac{\text{pos}}{10000^{2i / d_{\text{model}}}}\right) \end{aligned}
  • 直觉:类似于时钟系统,低维分量变化极快(像“秒针”),高维分量变化极慢(像“时针”),组合起来构成每个位置的唯一时间戳。
  • 缺点:直接加到 Embedding 上,随着层数加深会被网络变换逐渐稀释,且长文本外推能力差。

5.3 大模型霸主:RoPE(旋转位置编码)的几何魔术

RoPE(Rotary Position Embedding)由苏剑林等提出,是当今 LLaMA、Qwen、DeepSeek、Mistral 的统一绝对标准。 Ringi 导师解构:RoPE 旋转位置编码几何魔术小剧场

几何旋转本质

RoPE 不在输入 Embedding 上加向量,而是在计算 Attention 前,把 QQ 和 KK 向量的每相邻两个维度当成 2D 平面坐标,按位置角度进行逆时针旋转: [q2i′q2i+1′]=[cos⁡(pos⋅θi)−sin⁡(pos⋅θi)sin⁡(pos⋅θi)cos⁡(pos⋅θi)][q2iq2i+1]\begin{bmatrix} q'_{2i} \\ q'_{2i+1} \end{bmatrix} = \begin{bmatrix} \cos(\text{pos} \cdot \theta_i) & -\sin(\text{pos} \cdot \theta_i) \\ \sin(\text{pos} \cdot \theta_i) & \cos(\text{pos} \cdot \theta_i) \end{bmatrix} \begin{bmatrix} q_{2i} \\ q_{2i+1} \end{bmatrix}

神奇的数学特性:相对位置自然涌现

两个旋转后的向量做内积时,发生复数共轭相乘,绝对位置消去,结果仅依赖于它们的相对距离 (m−n)(m - n): ⟨RoPE(q,m),RoPE(k,n)⟩=g(q,k,m−n)\langle \text{RoPE}(q, m), \text{RoPE}(k, n) \rangle = g(q, k, m - n)
💡 Ringi 划重点:
  1. 为什么更好:语言中“第 3 个词与第 5 个词相距 2”比“它们分别在第 3 和第 5 位”重要得多。
  2. CUDA 算子融合:RoPE 是纯逐元素旋转,完全不需要分配显存矩阵,通常直接融合在 QKV GEMM 之后的同一个 CUDA Kernel 中完成!

6. LayerNorm 与残差连接——深层训练的定海神针

6.1 残差连接(Residual Connection):梯度高速公路

Ringi 导师解构:残差连接与 Pre-Norm 梯度高速公路 数学公式极简: y=x+SubLayer(x)y = x + \text{SubLayer}(x)

为什么要加这条捷径?

  1. 彻底解决梯度消失:反向传播求导时:
∂y∂x=1+∂SubLayer(x)∂x\frac{\partial y}{\partial x} = 1 + \frac{\partial \text{SubLayer}(x)}{\partial x} 因为始终存在一个恒等项 +1+1,梯度可以沿着主干道无损直达最浅层,使得堆叠 100 层以上的超深网络成为可能。 2. 增量学习:模型只需要学习每一层对输入的“微调增量”(Delta),学习难度大幅降低。

6.2 LayerNorm:特征维度的信号调节器

对单个 Token 向量在所有特征维度上做均值/方差归一化: LayerNorm(x)=x−μσ2+ϵ⊙γ+β\text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \odot \gamma + \beta 其中 μ\mu 为均值, σ2\sigma^2 为方差, γ,β\gamma, \beta 为可学习的缩放与平移参数。
🛠️ 现代演进:RMSNorm 现代大模型(LLaMA/Mistral 等)普遍改用 RMSNorm,省去计算均值 μ\mu 的步骤,直接用均方根归一化,效果几乎相同,但减少了一遍内存扫描,GPU 跑得更快! RMSNorm(x)=x1d∑i=1dxi2+ϵ⊙γ\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon}} \odot \gamma

6.3 Pre-Norm vs Post-Norm:数十亿美元训练稳定性之争


7. 完整的 Transformer Decoder Block 组装与参数手算

现在我们把所有零件拼装成一块完整的积木:

7.1 Pre-Norm Decoder Block 完整数据流与因果掩码

Ringi 导师解构:Transformer Decoder Block 完整流水线与残差高速公路

7.2 LLaMA-2-7B 全流程张量维度追踪(Shape 追踪表)

以大模型基准 LLaMA-2-7B 为例:
  • 序列长度 N=2048N = 2048,隐藏层维度 dmodel=4096d_{\text{model}} = 4096
  • 头数 h=32h = 32,单头维度 dk=128d_k = 128
  • FFN 中间维度 dff=11008d_{\text{ff}} = 11008
🌟 极其重要的性质: 输入是 (2048, 4096),输出依然是 (2048, 4096)! 形状完全一致,意味着这 32 层 Block 可以像乐高积木一样无缝串联。

7.3 参数量白板手算:7B 参数都花在了哪里?

1. 单个 Decoder Block 的参数量手算

  • Attention 部分:
WQ,WK,WV,WO  ⟹  4×(4096×4096)=4×16,777,216≈67.11 MW_Q, W_K, W_V, W_O \implies 4 \times (4096 \times 4096) = 4 \times 16,777,216 \approx \mathbf{67.11\text{ M}}
  • FFN (SwiGLU) 部分:
Wgate,Wup,Wdown  ⟹  3×(4096×11008)=3×45,088,768≈135.27 MW_{\text{gate}}, W_{\text{up}}, W_{\text{down}} \implies 3 \times (4096 \times 11008) = 3 \times 45,088,768 \approx \mathbf{135.27\text{ M}}
  • RMSNorm 缩放参数:
2×4096≈8.19 K2 \times 4096 \approx \mathbf{8.19\text{ K}}
  • 单层合计:
67.11M+135.27M≈202.38 M67.11\text{M} + 135.27\text{M} \approx \mathbf{202.38\text{ M}}

2. 全模型总参数量手算(32 层)

  • 32 层 Block: 32×202.38 M≈6,476 M32 \times 202.38\text{ M} \approx \mathbf{6,476\text{ M}}
  • Token Embedding: 32000×4096≈131.07 M32000 \times 4096 \approx \mathbf{131.07\text{ M}}
  • LM Head 输出头: 4096×32000≈131.07 M4096 \times 32000 \approx \mathbf{131.07\text{ M}}
  • 全模型精确总计:
6,476M+131M+131M≈6.738 B≈7B6,476\text{M} + 131\text{M} + 131\text{M} \approx \mathbf{6.738\text{ B}} \approx \mathbf{7\text{B}}

7.4 静态与动态显存规划公式(AI Infra 必背)

  1. 纯模型权重(FP16 / BF16,每个参数 2 字节):
6.74B×2 Bytes≈13.48 GB6.74\text{B} \times 2\text{ Bytes} \approx \mathbf{13.48\text{ GB}}
  1. 反向传播梯度(FP16,每个参数 2 字节):
6.74B×2 Bytes≈13.48 GB6.74\text{B} \times 2\text{ Bytes} \approx \mathbf{13.48\text{ GB}}
  1. AdamW 优化器状态(FP32 Master 权重 + 一阶动量 + 二阶动量 = 每个参数 16 字节):
6.74B×16 Bytes≈107.84 GB6.74\text{B} \times 16\text{ Bytes} \approx \mathbf{107.84\text{ GB}}
  1. 训练静态显存总计:
13.48+13.48+107.84=134.8 GB(≫80 GB!)13.48 + 13.48 + 107.84 = \mathbf{134.8\text{ GB}} \quad (\gg 80\text{ GB}!)
💡 这就是为什么必须使用 ZeRO / FSDP 显存切分技术,将优化器状态分摊到多张显卡上!

8. 从 Transformer 到 LLM:自回归生成与推理加速

8.1 自回归生成(Autoregressive):一次吐一个 Token

大模型推理是一个逐词生成的循环:

8.2 Prefill 与 Decode 两阶段:算力受限 vs 带宽受限

LLM 在线推理分为性质完全不同的两个阶段: Ringi 导师解构:LLM 自回归推理与 KV Cache 显存管理架构

8.3 KV Cache 的由来与计算量优化

在 Decode 阶段,生成第 nn 个词时:
  • 新 Token 的 QnQ_n 需要和历史前 (n−1)(n-1) 个 Token 的 K1,…,Kn−1K_1, \dots, K_{n-1} 做点积;
  • 然后与历史的 V1,…,Vn−1V_1, \dots, V_{n-1} 做加权求和。
  • 无 KV Cache:每生成一个新词,都要把前面所有的词重新算一遍 Q,K,VQ, K, V,计算复杂度为 O(N2⋅d2)O(N^2 \cdot d^2),算力极度浪费。
  • 开启 KV Cache:Prefill 阶段一次性将历史的 K,VK, V 缓存在 GPU 显存中,Decode 阶段每一步只需计算当前 1 个 Token 的 Q,K,VQ, K, V,并将新 K,VK, V 追加到 Cache 中,计算复杂度降为 O(N⋅d2)O(N \cdot d^2)!

8.4 KV Cache 显存开销与现代推理系统(PagedAttention)

1. KV Cache 显存占用手算公式(LLaMA-2-7B)

单 Token 显存=2×(层数 L)×(头数 h)×(头维度 dk)×(精度字节数)\text{单 Token 显存} = 2 \times (\text{层数 } L) \times (\text{头数 } h) \times (\text{头维度 } d_k) \times (\text{精度字节数}) 代入 LLaMA-2-7B(32 层,32 头,每头 128 维,FP16 占 2 字节): 单 Token 显存=2×32×32×128×2 Bytes=524,288 Bytes=512 KB\text{单 Token 显存} = 2 \times 32 \times 32 \times 128 \times 2\text{ Bytes} = \mathbf{524,288\text{ Bytes}} = \mathbf{512\text{ KB}}
  • 若一个请求长 4096 Token:
4096×512 KB=2 GB4096 \times 512\text{ KB} = \mathbf{2\text{ GB}}
  • 若并发 Batch Size 为 16:
16×2 GB=32 GB(已占据 80GB 显卡的近一半显存!)16 \times 2\text{ GB} = \mathbf{32\text{ GB}} \quad (\text{已占据 80GB 显卡的近一半显存!})
💡 Ringi 划重点: 传统的连续显存预分配会导致严重的内存碎片化(利用率通常 < 40\% )。 vLLM (PagedAttention) 借鉴了操作系统的虚拟内存分页机制,将 KV Cache 划分为固定大小的物理 Block,通过 Page Table 动态按需映射,彻底消除了显存碎片,将显存利用率提升至 96%96\% 以上!这也是现代大模型推理服务吞吐暴涨的核心秘诀。

9. 终极大复习与极客速记口诀


10. Google 大神 × Ringi 检验清单与进阶路线

10.1 小白通关自我检验清单

完成本文学习后,请尝试合上笔记回答以下 10 个灵魂拷问:
  • Q1:能不看资料,在纸上画出一个完整的 Pre-Norm Decoder Block 结构流向图,并标注所有残差与 Norm 位置吗?
  • Q2:能说清为什么现代大模型全都在用 Decoder-only 架构,而不是 Encoder-Decoder 吗?
  • Q3:能用自己的话说清 Q,K,VQ, K, V 三个矩阵的物理含义,以及点积的几何直觉吗?
  • Q4:能默写 Attention 核心公式,并准确解释为什么要除以 dk\sqrt{d_k} 吗?
  • Q5:能推导出 Self-Attention 的 O(N2)O(N^2) 计算量与显存开销,并说明 FlashAttention 优化的核心思想吗?
  • Q6:能说明为什么 Multi-Head Attention 天然适合做 Megatron-LM 张量并行(TP)切分吗?
  • Q7:能独立手算出 LLaMA-2-7B 的单层 Block 和全模型参数量(~6.7B),并解释为什么 FFN 占了 67% 吗?
  • Q8:能解释为什么 RoPE 旋转位置编码做内积时只与相对位置差 (m−n)(m-n) 有关吗?
  • Q9:能说清 Prefill 和 Decode 两个阶段的物理特性区别(Compute Bound vs Memory Bound)吗?
  • Q10:能手算出给定配置下 KV Cache 的显存大小,并解释 vLLM 的 PagedAttention 解决了什么痛点吗?

10.2 进阶推荐阅读与论文清单

必读经典论文(AI Infra 基石)

  1. Transformer 开山之作:Vaswani et al., 2017. Attention Is All You Need
  2. RoPE 旋转位置编码:Su et al., 2021. RoFormer: Enhanced Transformer with Rotary Position Embedding
  3. SwiGLU 门控激活:Shazeer, 2020. GLU Variants Improve Transformer
  4. LLaMA 架构标杆:Touvron et al., 2023. LLaMA: Open and Efficient Foundation Language Models
  5. FlashAttention 显存优化:Dao et al., 2022. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
  6. GQA 分组注意力:Ainslie et al., 2023. GQA: Training Generalized Multi-Query Transformer Models
  7. vLLM PagedAttention:Kwon et al., 2023. Efficient Memory Management for Large Language Model Serving with PagedAttention
  8. Transformer架构-快速入门篇:Transformer架构-快速入门篇

经典视频与实践教程

  • Andrej Karpathy:《Let’s build GPT from scratch》(YouTube 必看代码级实战)
  • 3Blue1Brown:《But what is a GPT? Visual Intro to Transformers》(几何直觉级神作)
  • Jay Alammar:《The Illustrated Transformer》(图解 Transformer 经典图文)